/* ----------------------------------------------------------------------------
Copyright (c) 2018-2021, Microsoft Research, Daan Leijen
This is free software; you can redistribute it and/or modify it under the
terms of the MIT license. A copy of the license can be found in the file
"LICENSE" at the root of this distribution.
-----------------------------------------------------------------------------*/

#include "mimalloc.h"
#include "mimalloc/internal.h"
#include "mimalloc/prim.h"  // mi_prim_get_default_heap

#include <string.h>     // memset

// ------------------------------------------------------
// Aligned Allocation
// ------------------------------------------------------

// Fallback primitive aligned allocation -- split out for better codegen
/* Fallback primitive aligned allocation.
Allocates memory of the given size and alignment, with an optional offset,
from the provided heap. The allocated memory is zeroed if the zero parameter is
true. Returns a pointer to the allocated memory, or NULL if the allocation
fails.
*/
static mi_decl_noinline void *mi_heap_malloc_zero_aligned_at_fallback(
	mi_heap_t *const heap, const size_t size, const size_t alignment,
	const size_t offset, const bool zero) mi_attr_noexcept
{
  mi_assert_internal(size <= PTRDIFF_MAX);
  mi_assert_internal(alignment != 0 && _mi_is_power_of_two(alignment));

  const uintptr_t align_mask = alignment - 1;  // for any x, `(x & align_mask) == (x % alignment)`
  const size_t padsize = size + MI_PADDING_SIZE;

  // use regular allocation if it is guaranteed to fit the alignment constraints
  if (offset==0 && alignment<=padsize && padsize<=MI_MAX_ALIGN_GUARANTEE && (padsize&align_mask)==0) {
    void* p = _mi_heap_malloc_zero(heap, size, zero);
    mi_assert_internal(p == NULL || ((uintptr_t)p % alignment) == 0);
    return p;
  }

  void* p;
  size_t oversize;
  if mi_unlikely(alignment > MI_ALIGNMENT_MAX) {
    // use OS allocation for very large alignment and allocate inside a huge page (dedicated segment with 1 page)
    // This can support alignments >= MI_SEGMENT_SIZE by ensuring the object can be aligned at a point in the
    // first (and single) page such that the segment info is `MI_SEGMENT_SIZE` bytes before it (so it can be found by aligning the pointer down)
    if mi_unlikely(offset != 0) {
      // todo: cannot support offset alignment for very large alignments yet
      #if MI_DEBUG > 0
      _mi_error_message(EOVERFLOW, "aligned allocation with a very large alignment cannot be used with an alignment offset (size %zu, alignment %zu, offset %zu)\n", size, alignment, offset);
      #endif
      return NULL;
    }
    oversize = (size <= MI_SMALL_SIZE_MAX ? MI_SMALL_SIZE_MAX + 1 /* ensure we use generic malloc path */ : size);
    p = _mi_heap_malloc_zero_ex(heap, oversize, false, alignment); // the page block size should be large enough to align in the single huge page block
    // zero afterwards as only the area from the aligned_p may be committed!
    if (p == NULL) return NULL;
  }
  else {
    // otherwise over-allocate
    oversize = size + alignment - 1;
    p = _mi_heap_malloc_zero(heap, oversize, zero);
    if (p == NULL) return NULL;
  }

  // .. and align within the allocation
  const uintptr_t poffset = ((uintptr_t)p + offset) & align_mask;
  const uintptr_t adjust  = (poffset == 0 ? 0 : alignment - poffset);
  mi_assert_internal(adjust < alignment);
  void* aligned_p = (void*)((uintptr_t)p + adjust);
  if (aligned_p != p) {
    mi_page_t* page = _mi_ptr_page(p);
    mi_page_set_has_aligned(page, true);
    _mi_padding_shrink(page, (mi_block_t*)p, adjust + size);
  }
  // todo: expand padding if overallocated ?

  mi_assert_internal(mi_page_usable_block_size(_mi_ptr_page(p)) >= adjust + size);
  mi_assert_internal(p == _mi_page_ptr_unalign(_mi_ptr_segment(aligned_p), _mi_ptr_page(aligned_p), aligned_p));
  mi_assert_internal(((uintptr_t)aligned_p + offset) % alignment == 0);
  mi_assert_internal(mi_usable_size(aligned_p)>=size);
  mi_assert_internal(mi_usable_size(p) == mi_usable_size(aligned_p)+adjust);

  // now zero the block if needed
  if (alignment > MI_ALIGNMENT_MAX) {
    // for the tracker, on huge aligned allocations only from the start of the large block is defined
    mi_track_mem_undefined(aligned_p, size);
    if (zero) {
      _mi_memzero_aligned(aligned_p, mi_usable_size(aligned_p));
    }
  }

  if (p != aligned_p) {
    mi_track_align(p,aligned_p,adjust,mi_usable_size(aligned_p));
  }
  return aligned_p;
}

// Primitive aligned allocation
/**
 * Allocates `size` bytes of memory aligned on a boundary of
 * `alignment` at the specified `offset` from the provided `heap`.
 * The allocated memory is zeroed if `zero` is true.
 * Returns a pointer to the allocated memory, or `NULL` if the request fails.
 * This is an internal implementation function.
 */
static void *mi_heap_malloc_zero_aligned_at(mi_heap_t *const heap,
					    const size_t size,
					    const size_t alignment,
					    const size_t offset,
					    const bool zero) mi_attr_noexcept
{
  // note: we don't require `size > offset`, we just guarantee that the address at offset is aligned regardless of the allocated size.
  if mi_unlikely(alignment == 0 || !_mi_is_power_of_two(alignment)) { // require power-of-two (see <https://en.cppreference.com/w/c/memory/aligned_alloc>)
    #if MI_DEBUG > 0
    _mi_error_message(EOVERFLOW, "aligned allocation requires the alignment to be a power-of-two (size %zu, alignment %zu)\n", size, alignment);
    #endif
    return NULL;
  }

  if mi_unlikely(size > PTRDIFF_MAX) {          // we don't allocate more than PTRDIFF_MAX (see <https://sourceware.org/ml/libc-announce/2019/msg00001.html>)
    #if MI_DEBUG > 0
    _mi_error_message(EOVERFLOW, "aligned allocation request is too large (size %zu, alignment %zu)\n", size, alignment);
    #endif
    return NULL;
  }
  const uintptr_t align_mask = alignment-1;       // for any x, `(x & align_mask) == (x % alignment)`
  const size_t padsize = size + MI_PADDING_SIZE;  // note: cannot overflow due to earlier size > PTRDIFF_MAX check

  // try first if there happens to be a small block available with just the right alignment
  if mi_likely(padsize <= MI_SMALL_SIZE_MAX && alignment <= padsize) {
    mi_page_t* page = _mi_heap_get_free_small_page(heap, padsize);
    const bool is_aligned = (((uintptr_t)page->free+offset) & align_mask)==0;
    if mi_likely(page->free != NULL && is_aligned)
    {
      #if MI_STAT>1
      mi_heap_stat_increase(heap, malloc, size);
      #endif
      void* p = _mi_page_malloc(heap, page, padsize, zero); // TODO: inline _mi_page_malloc
      mi_assert_internal(p != NULL);
      mi_assert_internal(((uintptr_t)p + offset) % alignment == 0);
      mi_track_malloc(p,size,zero);
      return p;
    }
  }
  // fallback
  return mi_heap_malloc_zero_aligned_at_fallback(heap, size, alignment, offset, zero);
}


// ------------------------------------------------------
// Optimized mi_heap_malloc_aligned / mi_malloc_aligned
// ------------------------------------------------------

/**
 * Allocates `size` bytes of memory aligned on a boundary of
 * `alignment` at the specified `offset`. Returns a pointer to the
 * allocated memory, or `NULL` if the request fails. Uses the provided
 * `heap`.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_heap_malloc_aligned_at(mi_heap_t *heap, size_t size, size_t alignment,
			  size_t offset) mi_attr_noexcept
{
  return mi_heap_malloc_zero_aligned_at(heap, size, alignment, offset, false);
}

/**
 * Allocates `size` bytes of memory aligned on a boundary of
 * `alignment`. Returns a pointer to the allocated memory, or `NULL` if
 * the request fails. Uses the provided `heap`.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_heap_malloc_aligned(mi_heap_t *heap, size_t size,
		       size_t alignment) mi_attr_noexcept
{
  if mi_unlikely(alignment == 0 || !_mi_is_power_of_two(alignment)) return NULL;
  #if !MI_PADDING
  // without padding, any small sized allocation is naturally aligned (see also `_mi_segment_page_start`)
  if mi_likely(_mi_is_power_of_two(size) && size >= alignment && size <= MI_SMALL_SIZE_MAX)
  #else
  // with padding, we can only guarantee this for fixed alignments
  if mi_likely((alignment == sizeof(void*) || (alignment == MI_MAX_ALIGN_SIZE && size > (MI_MAX_ALIGN_SIZE/2)))
		&& size <= MI_SMALL_SIZE_MAX)
  #endif
  {
    // fast path for common alignment and size
    return mi_heap_malloc_small(heap, size);
  }
  else {
    return mi_heap_malloc_aligned_at(heap, size, alignment, 0);
  }
}

// ensure a definition is emitted
#if defined(__cplusplus)
static void* _mi_heap_malloc_aligned = (void*)&mi_heap_malloc_aligned;
#endif

// ------------------------------------------------------
// Aligned Allocation
// ------------------------------------------------------

/**
 * Allocates `size` bytes of memory aligned on a boundary of
 * `alignment` at the specified `offset`. The memory is initialized
 * to zero. Returns a pointer to the allocated memory, or `NULL` if
 * the request fails. Uses the provided `heap`.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_heap_zalloc_aligned_at(mi_heap_t *heap, size_t size, size_t alignment,
			  size_t offset) mi_attr_noexcept
{
  return mi_heap_malloc_zero_aligned_at(heap, size, alignment, offset, true);
}

/**
 * Allocates `size` bytes of memory aligned on a boundary of
 * `alignment`. The memory is initialized to zero.
 * Returns a pointer to the allocated memory, or `NULL` if the
 * request fails. Uses the provided `heap`.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_heap_zalloc_aligned(mi_heap_t *heap, size_t size,
		       size_t alignment) mi_attr_noexcept
{
  return mi_heap_zalloc_aligned_at(heap, size, alignment, 0);
}

/**
 * Allocates memory for an array of `count` elements of `size` bytes each.
 * The allocated memory is initialized to zero and aligned on a boundary of
 * `alignment` at the specified `offset`. Returns a pointer to the allocated
 * memory, or `NULL` if the request fails. Uses the provided `heap`.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_heap_calloc_aligned_at(mi_heap_t *heap, size_t count, size_t size,
			  size_t alignment, size_t offset) mi_attr_noexcept
{
  size_t total;
  if (mi_count_size_overflow(count, size, &total)) return NULL;
  return mi_heap_zalloc_aligned_at(heap, total, alignment, offset);
}

/**
 * Allocates memory for an array of `count` elements of `size` bytes each.
 * The allocated memory is initialized to zero and aligned on a boundary of
 * `alignment`. Returns a pointer to the allocated memory, or `NULL` if the
 * request fails. Uses the provided `heap`.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_heap_calloc_aligned(mi_heap_t *heap, size_t count, size_t size,
		       size_t alignment) mi_attr_noexcept
{
  return mi_heap_calloc_aligned_at(heap,count,size,alignment,0);
}

/**
 * Allocates `size` bytes of memory aligned on a boundary of `alignment`
 * at the specified `offset`. Returns a pointer to the allocated memory,
 * or `NULL` if the request fails. Uses the provided heap.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_malloc_aligned_at(size_t size, size_t alignment,
		     size_t offset) mi_attr_noexcept
{
  return mi_heap_malloc_aligned_at(mi_prim_get_default_heap(), size, alignment, offset);
}

/**
 * Allocates `size` bytes of memory aligned on a boundary of `alignment`.
 * Returns a pointer to the allocated memory, or `NULL` if the request fails.
 * Uses the default heap.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_malloc_aligned(size_t size, size_t alignment) mi_attr_noexcept
{
  return mi_heap_malloc_aligned(mi_prim_get_default_heap(), size, alignment);
}

/**
 * Allocates `size` bytes of memory aligned on a boundary of `alignment`
 * at the specified `offset`. The allocated memory is initialized to
 * zero. Returns a pointer to the allocated memory, or `NULL` if the
 * request fails. Uses the provided heap.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_zalloc_aligned_at(size_t size, size_t alignment,
		     size_t offset) mi_attr_noexcept
{
  return mi_heap_zalloc_aligned_at(mi_prim_get_default_heap(), size, alignment, offset);
}

/**
 * Allocates `size` bytes of memory aligned on a boundary of `alignment`.
 * The allocated memory is initialized to zero.
 * Returns a pointer to the allocated memory, or `NULL` if the request fails.
 * Uses the default heap.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_zalloc_aligned(size_t size, size_t alignment) mi_attr_noexcept
{
  return mi_heap_zalloc_aligned(mi_prim_get_default_heap(), size, alignment);
}

/**
 * Allocates memory for an array of `count` elements of size `size`.
 * The allocated memory is aligned by `alignment` at the specified `offset`.
 * Returns a pointer to the allocated memory, or `NULL` if the request fails.
 * Uses the provided heap.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_calloc_aligned_at(size_t count, size_t size, size_t alignment,
		     size_t offset) mi_attr_noexcept
{
  return mi_heap_calloc_aligned_at(mi_prim_get_default_heap(), count, size, alignment, offset);
}

/**
 * Allocates memory for an array of `count` elements of size `size`.
 * The allocated memory is aligned by `alignment`.
 * Returns a pointer to the allocated memory, or `NULL` if the request fails.
 * Uses the default heap.
 */
mi_decl_nodiscard mi_decl_restrict void *
mi_calloc_aligned(size_t count, size_t size, size_t alignment) mi_attr_noexcept
{
  return mi_heap_calloc_aligned(mi_prim_get_default_heap(), count, size, alignment);
}


// ------------------------------------------------------
// Aligned re-allocation
// ------------------------------------------------------

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment` at the specified `offset`.
 * If `zero` is true, reallocated memory is zeroed.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate.
 * Uses provided `heap`.
 */
static void *mi_heap_realloc_zero_aligned_at(mi_heap_t *heap, void *p,
					     size_t newsize, size_t alignment,
					     size_t offset,
					     bool zero) mi_attr_noexcept
{
  mi_assert(alignment > 0);
  if (alignment <= sizeof(uintptr_t)) return _mi_heap_realloc_zero(heap,p,newsize,zero);
  if (p == NULL) return mi_heap_malloc_zero_aligned_at(heap,newsize,alignment,offset,zero);
  size_t size = mi_usable_size(p);
  if (newsize <= size && newsize >= (size - (size / 2))
      && (((uintptr_t)p + offset) % alignment) == 0) {
    return p;  // reallocation still fits, is aligned and not more than 50% waste
  }
  else {
    // note: we don't zero allocate upfront so we only zero initialize the expanded part
    void* newp = mi_heap_malloc_aligned_at(heap,newsize,alignment,offset);
    if (newp != NULL) {
      if (zero && newsize > size) {
	// also set last word in the previous allocation to zero to ensure any padding is zero-initialized
	size_t start = (size >= sizeof(intptr_t) ? size - sizeof(intptr_t) : 0);
	_mi_memzero((uint8_t*)newp + start, newsize - start);
      }
      _mi_memcpy_aligned(newp, p, (newsize > size ? size : newsize));
      mi_free(p); // only free if successful
    }
    return newp;
  }
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment`.
 * If `zero` is true, reallocated memory is zeroed.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate.
 * Uses provided `heap`.
 */
static void *mi_heap_realloc_zero_aligned(mi_heap_t *heap, void *p,
					  size_t newsize, size_t alignment,
					  bool zero) mi_attr_noexcept
{
  mi_assert(alignment > 0);
  if (alignment <= sizeof(uintptr_t)) return _mi_heap_realloc_zero(heap,p,newsize,zero);
  size_t offset = ((uintptr_t)p % alignment); // use offset of previous allocation (p can be NULL)
  return mi_heap_realloc_zero_aligned_at(heap,p,newsize,alignment,offset,zero);
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment` at the specified `offset`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate.
 * Uses provided `heap`.
 */
mi_decl_nodiscard void *
mi_heap_realloc_aligned_at(mi_heap_t *heap, void *p, size_t newsize,
			   size_t alignment, size_t offset) mi_attr_noexcept
{
  return mi_heap_realloc_zero_aligned_at(heap,p,newsize,alignment,offset,false);
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate.
 * Uses provided `heap`.
 */
mi_decl_nodiscard void *
mi_heap_realloc_aligned(mi_heap_t *heap, void *p, size_t newsize,
			size_t alignment) mi_attr_noexcept
{
  return mi_heap_realloc_zero_aligned(heap,p,newsize,alignment,false);
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment` at the specified `offset`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate.
 * Uses provided `heap`.
 */
mi_decl_nodiscard void *
mi_heap_rezalloc_aligned_at(mi_heap_t *heap, void *p, size_t newsize,
			    size_t alignment, size_t offset) mi_attr_noexcept
{
  return mi_heap_realloc_zero_aligned_at(heap, p, newsize, alignment, offset, true);
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate.
 * Uses provided `heap`.
 */
mi_decl_nodiscard void *
mi_heap_rezalloc_aligned(mi_heap_t *heap, void *p, size_t newsize,
			 size_t alignment) mi_attr_noexcept
{
  return mi_heap_realloc_zero_aligned(heap, p, newsize, alignment, true);
}

/**
 * Reallocates memory block `p` to hold space for `newcount` elements each
 * of size `size`. Memory is aligned by `alignment` at the specified
 * `offset`. Returns pointer to reallocated memory of adequate size to
 * hold `newcount` elements of size `size`, or `NULL` if unable to allocate
 * or reallocate. Uses provided `heap`.
 */
mi_decl_nodiscard void *
mi_heap_recalloc_aligned_at(mi_heap_t *heap, void *p, size_t newcount,
			    size_t size, size_t alignment,
			    size_t offset) mi_attr_noexcept
{
  size_t total;
  if (mi_count_size_overflow(newcount, size, &total)) return NULL;
  return mi_heap_rezalloc_aligned_at(heap, p, total, alignment, offset);
}

/**
 * Reallocates memory block `p` to hold space for `newcount` elements each
 * of size `size`. Memory is aligned by `alignment`.
 * Returns pointer to reallocated memory of adequate size to hold `newcount`
 * elements of size `size`, or `NULL` if unable to allocate or reallocate.
 * Uses provided `heap`.
 */
mi_decl_nodiscard void *
mi_heap_recalloc_aligned(mi_heap_t *heap, void *p, size_t newcount, size_t size,
			 size_t alignment) mi_attr_noexcept
{
  size_t total;
  if (mi_count_size_overflow(newcount, size, &total)) return NULL;
  return mi_heap_rezalloc_aligned(heap, p, total, alignment);
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment` at the specified `offset`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate.
 * Uses default heap.
 */
mi_decl_nodiscard void *mi_realloc_aligned_at(void *p, size_t newsize,
					      size_t alignment,
					      size_t offset) mi_attr_noexcept
{
  return mi_heap_realloc_aligned_at(mi_prim_get_default_heap(), p, newsize, alignment, offset);
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate.
 * Uses default heap.
 */
mi_decl_nodiscard void *mi_realloc_aligned(void *p, size_t newsize,
					   size_t alignment) mi_attr_noexcept
{
  return mi_heap_realloc_aligned(mi_prim_get_default_heap(), p, newsize, alignment);
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment` at the specified `offset`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate. Zeroes new allocation if `p` is `NULL`.
 * Uses default heap.
 */
mi_decl_nodiscard void *mi_rezalloc_aligned_at(void *p, size_t newsize,
					       size_t alignment,
					       size_t offset) mi_attr_noexcept
{
  return mi_heap_rezalloc_aligned_at(mi_prim_get_default_heap(), p, newsize, alignment, offset);
}

/**
 * Reallocates memory block `p` to new size `newsize`.
 * Memory is aligned by `alignment`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate. Zeroes new allocation if `p` is `NULL`.
 * Uses default heap.
 */
mi_decl_nodiscard void *mi_rezalloc_aligned(void *p, size_t newsize,
					    size_t alignment) mi_attr_noexcept
{
  return mi_heap_rezalloc_aligned(mi_prim_get_default_heap(), p, newsize, alignment);
}

/**
 * Reallocates memory block `p` to new size capable of holding `newcount`
 * elements of `size` bytes each. Memory is aligned by `alignment` at
 * the specified `offset`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate. Reuses existing allocation if sufficient size
 * or zeroes new allocation if `p` is `NULL`. Uses default heap.
 */
mi_decl_nodiscard void *mi_recalloc_aligned_at(void *p, size_t newcount,
					       size_t size, size_t alignment,
					       size_t offset) mi_attr_noexcept
{
  return mi_heap_recalloc_aligned_at(mi_prim_get_default_heap(), p, newcount, size, alignment, offset);
}

/**
 * Reallocates memory block `p` to new size capable of holding `newcount`
 * elements of `size` bytes each. Memory is aligned by `alignment`.
 * Returns pointer to reallocated memory of adequate size, or `NULL` if unable
 * to allocate or reallocate. Reuses existing allocation if sufficient size
 * or zeroes new allocation if `p` is `NULL`. Uses default heap.
 */
mi_decl_nodiscard void *mi_recalloc_aligned(void *p, size_t newcount,
					    size_t size,
					    size_t alignment) mi_attr_noexcept
{
  return mi_heap_recalloc_aligned(mi_prim_get_default_heap(), p, newcount, size, alignment);
}
